package com.zlx.mall.basic.trace;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

/**
 * @author qingke
 **/
public class MDCExecutorWrapper {

    public static <V> Callable<V> wrap(final Callable<V> callable, final Map<String, String> contextMap) {
        return () -> {
            if (null == contextMap) {
                MDC.clear();
            } else {
                MDC.setContextMap(contextMap);
            }
            MDCUtil.getOrCreateTraceId();
            try {
                return callable.call();
            } finally {
                MDC.clear();
            }
        };
    }

    public static Runnable wrap(final Runnable runnable, final Map<String, String> contextMap) {
        return () -> {
            if (null == contextMap) {
                MDC.clear();
            } else {
                MDC.setContextMap(contextMap);
            }
            MDCUtil.getOrCreateTraceId();
            try {
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }
}
